## 删除标签文件中除了"USV"外的其他物项的object
import os
import shutil
import xml.etree.ElementTree as ET

def filter_and_copy_labels(source_folder, target_folder, keep_name="USV"):
    annotations_source = os.path.join(source_folder, "Annotations")
    images_source = os.path.join(source_folder, "JPEGImages")
    annotations_target = os.path.join(target_folder, "Annotations")
    images_target = os.path.join(target_folder, "JPEGImages")
    
    os.makedirs(annotations_target, exist_ok=True)
    os.makedirs(images_target, exist_ok=True)
    
    for filename in os.listdir(annotations_source):
        if filename.endswith(".xml"):
            xml_path = os.path.join(annotations_source, filename)
            tree = ET.parse(xml_path)
            root = tree.getroot()
            
            # 找到所有 object 标签并移除 name 不是 USV 的对象
            for obj in root.findall("object"):
                name = obj.find("name").text
                if name != keep_name:
                    root.remove(obj)
            
            # 只有当至少有一个 USV 对象时才复制文件
            if root.findall("object"):
                tree.write(os.path.join(annotations_target, filename), encoding="utf-8")
                
                # 复制对应的图片文件
                img_filename = os.path.splitext(filename)[0] + ".jpg"
                img_path = os.path.join(images_source, img_filename)
                if os.path.exists(img_path):
                    shutil.copy(img_path, os.path.join(images_target, img_filename))
                else:
                    print(f"Warning: Corresponding image {img_filename} not found.")
    
    print("Processing completed!")

# 示例使用，替换为你的实际路径
source_folder = "/home/shuai/Documents/mydataset/my_usv_dock_exp/VOCdevkit_USV_exp"
target_folder = "/home/shuai/Documents/mydataset/my_usv_dock_exp/VOCdevkit_onlyUSV_exp"
filter_and_copy_labels(source_folder, target_folder)
